import pandas as pd
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np


class MovieRecommender:
    def __init__(self, db):
        self.db = db
        self.df = None
        self.all_regions = set()
        self.all_types = set()
        self.mlb_region = None
        self.mlb_type = None
        self.combined_matrix = None

    def preprocess_data(self):
        # 使用 pandas 查询数据库
        self.df = pd.read_sql("SELECT id, title, region, type FROM movies", self.db.engine)
        # 将 region 列转换为每个电影的国家列表
        self.df['regions'] = self.df['region'].str.split(' ')
        # 将 types 列拆分为类型列表
        self.df['types'] = self.df['type'].str.split(' ')
        # 获取所有不重复的国家和类型
        self.all_regions = set([region for regions_list in self.df['regions'] for region in regions_list])
        self.all_types = set([type for types_list in self.df['types'] if types_list for type in types_list])

        # 独热编码
        self.mlb_region = MultiLabelBinarizer()
        self.mlb_type = MultiLabelBinarizer()

        region_matrix = self.mlb_region.fit_transform(self.df['regions'])
        type_matrix = self.mlb_type.fit_transform(self.df['types'])

        # 合并国家和类型的独热编码矩阵
        self.combined_matrix = np.hstack((region_matrix, type_matrix))

    def rec_by_region_and_type(self, input_region, input_type, count=10):
        if self.df is not None and self.mlb_region is not None and self.mlb_type is not None:
            # 为输入生成独热编码向量
            input_region = input_region.split(" ")
            input_type = input_type.split(" ")
            input_region_vec = self.mlb_region.transform(np.reshape(input_region, (1, -1)))
            input_type_vec = self.mlb_type.transform(np.reshape(input_type, (1, -1)))
            input_vec = np.hstack((input_region_vec, input_type_vec))

            # 计算相似度
            similarities = cosine_similarity(self.combined_matrix, input_vec)

            # 获取最相似的count部电影
            topN_indices = similarities.reshape((-1,)).argsort()[-count:][::-1]
            recommended_movies_id = self.df['id'].iloc[topN_indices]

            return recommended_movies_id.tolist()

    def rec_by_movie(self, input_movie, count=10):
        # 检查输入电影是否在数据库中
        if input_movie not in self.df['title'].values:
            return "输入的电影不在数据库中，请输入有效的电影名。"

        # 获取输入电影的特征向量
        input_movie_index = self.df[self.df['title'] == input_movie].index[0]
        input_movie_vec = self.combined_matrix[input_movie_index]

        # 计算相似度
        similarities = cosine_similarity(self.combined_matrix, input_movie_vec.reshape(1, -1))

        # 获取最相似的count部电影（排除输入电影本身）
        topN_indices = similarities.reshape((-1,)).argsort()[-(count + 1):][::-1]  # 取前n+1个，因为第1个是电影本身
        topN_indices = topN_indices[topN_indices != input_movie_index]  # 排除电影本身
        recommended_movies_id = self.df['id'].iloc[topN_indices]

        return recommended_movies_id.tolist()
